{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a94de9a7",
   "metadata": {
    "colab_type": "text",
    "id": "EgiF12Hf1Dhs"
   },
   "source": [
    "This notebook provides examples to go along with the [textbook](http://manipulation.csail.mit.edu/robot.html).  I recommend having both windows open, side-by-side!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46d55a45",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eeMrMI0-1Dhu",
    "lines_to_end_of_cell_marker": 2
   },
   "outputs": [],
   "source": [
    "from pydrake.all import ModelVisualizer, Simulator, StartMeshcat\n",
    "\n",
    "from manipulation import ConfigureParser, running_as_notebook\n",
    "from manipulation.station import MakeHardwareStation, load_scenario"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "388ddbad",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41462d03",
   "metadata": {},
   "source": [
    "# PR2 model example\n",
    "\n",
    "First we'll use the ModelVisualizer to inspect the model (note that [the sliders in ModelVisualizer do not currently respect the kinematic \"mimic\" joint in the fingers](https://github.com/RobotLocomotion/drake/issues/18917))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bce34ac",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5SjOClhTltPk"
   },
   "outputs": [],
   "source": [
    "visualizer = ModelVisualizer(meshcat=meshcat)\n",
    "ConfigureParser(visualizer.parser())\n",
    "visualizer.AddModels(\n",
    "    url=\"package://drake/examples/pr2/models/pr2_description/urdf/pr2_simplified.urdf\"\n",
    ")\n",
    "visualizer.Run(loop_once=not running_as_notebook)\n",
    "meshcat.DeleteAddedControls()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2314f6eb",
   "metadata": {},
   "source": [
    "Now we can use HardwareStation to create a basic simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8610a38",
   "metadata": {},
   "outputs": [],
   "source": [
    "scenario_data = \"\"\"\n",
    "directives:\n",
    "- add_model:\n",
    "    name: pr2\n",
    "    file: package://drake/examples/pr2/models/pr2_description/urdf/pr2_simplified.urdf\n",
    "model_drivers:\n",
    "    pr2: !JointStiffnessDriver\n",
    "        gains:\n",
    "            # TODO(russt): tune these!\n",
    "            x_motor:\n",
    "                kp: 600\n",
    "                kd: 120\n",
    "            y_motor:\n",
    "                kp: 600\n",
    "                kd: 120\n",
    "            theta_motor:\n",
    "                kp: 600\n",
    "                kd: 120\n",
    "            torso_lift_motor:\n",
    "                kp: 600\n",
    "                kd: 120\n",
    "            head_pan_motor:\n",
    "                kp: 100\n",
    "                kd: 20\n",
    "            head_tilt_motor:\n",
    "                kp: 100\n",
    "                kd: 20\n",
    "            r_upper_arm_roll_motor:\n",
    "                kp: 600\n",
    "                kd: 120\n",
    "            r_shoulder_pan_motor:\n",
    "                kp: 600\n",
    "                kd: 120\n",
    "            r_shoulder_lift_motor:\n",
    "                kp: 600\n",
    "                kd: 120\n",
    "            r_forearm_roll_motor:\n",
    "                kp: 400\n",
    "                kd: 80\n",
    "            r_elbow_flex_motor:\n",
    "                kp: 400\n",
    "                kd: 80\n",
    "            r_wrist_flex_motor:\n",
    "                kp: 200\n",
    "                kd: 40\n",
    "            r_wrist_roll_motor:\n",
    "                kp: 200\n",
    "                kd: 40\n",
    "            r_gripper_l_finger_motor:\n",
    "                kp: 100\n",
    "                kd: 20\n",
    "            l_upper_arm_roll_motor:\n",
    "                kp: 600\n",
    "                kd: 120\n",
    "            l_shoulder_pan_motor:\n",
    "                kp: 600\n",
    "                kd: 120\n",
    "            l_shoulder_lift_motor:\n",
    "                kp: 600\n",
    "                kd: 120\n",
    "            l_forearm_roll_motor:\n",
    "                kp: 400\n",
    "                kd: 80\n",
    "            l_elbow_flex_motor:\n",
    "                kp: 400\n",
    "                kd: 80\n",
    "            l_wrist_flex_motor:\n",
    "                kp: 200\n",
    "                kd: 40\n",
    "            l_wrist_roll_motor:\n",
    "                kp: 200\n",
    "                kd: 40\n",
    "            l_gripper_l_finger_motor:\n",
    "                kp: 100\n",
    "                kd: 20\n",
    "\"\"\"\n",
    "\n",
    "scenario = load_scenario(data=scenario_data)\n",
    "station = MakeHardwareStation(scenario, meshcat)\n",
    "simulator = Simulator(station)\n",
    "context = simulator.get_mutable_context()\n",
    "\n",
    "plant = station.GetSubsystemByName(\"plant\")\n",
    "pr2 = plant.GetModelInstanceByName(\"pr2\")\n",
    "\n",
    "mimic_joints = [\n",
    "    \"gripper_r_finger_joint\",\n",
    "    \"gripper_l_finger_tip_joint\",\n",
    "    \"gripper_r_finger_tip_joint\",\n",
    "    \"gripper_r_finger_joint\",\n",
    "    \"gripper_l_finger_tip_joint\",\n",
    "    \"gripper_r_finger_tip_joint\",\n",
    "]\n",
    "x0 = station.GetOutputPort(\"pr2.state_estimated\").Eval(context)\n",
    "x0_wout_mimics = []\n",
    "for i, state_name in enumerate(plant.GetStateNames(pr2)):\n",
    "    if all(mimic not in state_name for mimic in mimic_joints):\n",
    "        x0_wout_mimics.append(x0[i])\n",
    "\n",
    "station.GetInputPort(\"pr2.desired_state\").FixValue(context, x0_wout_mimics)\n",
    "simulator.AdvanceTo(0.1)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.8 64-bit",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
